#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: yerunyuan@163.com
"""

import argparse
import os
import sys
import shutil
import re

sys.path.append('.')

"""
   拷贝图片
"""


def main():
    parser = argparse.ArgumentParser(description="Copy Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/DataSet/FaceVideo/images", help="path to image dir", type=str
    )
    parser.add_argument(
        "--bbox_count_file", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin/bbox_count_list.txt", help="path to file bbox count list", type=str
    )
    parser.add_argument(
        "--bbox_file", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin/bbox_list.txt", help="path to file bbox list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/FaceVideo", help="path to output", type=str
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    os.system('chmod a+wr {}'.format(output_dir))

    pattern = re.compile(r'F([-\d]+)B([-\d]+)')

    out_file_path = os.path.join(output_dir, 'file_bbox_list.txt')
    out_file = open(out_file_path, 'w')

    out_img_dir = os.path.join(output_dir, 'images')
    if not os.path.exists(out_img_dir):
        os.makedirs(out_img_dir)
    os.system('chmod a+wr {}'.format(out_img_dir))

    file_count = 0
    with open(args.bbox_count_file, 'r') as bbox_count_file:
        with open(args.bbox_file, 'r') as bbox_file:
            for line in bbox_count_file.readlines():
                lines = line.strip().split(',')
                file_dir = os.path.splitext(lines[0])[0]
                bbox_count = int(lines[1])
                for _ in range(bbox_count):
                    bbox_infos = bbox_file.readline()
                    bbox_infos = bbox_infos.strip().split()
                    file_name = bbox_infos[0]
                    # bbox = list(map(float, bbox_infos[2].split(',')))
                    bbox = bbox_infos[2]

                    fr_num, b_num = map(int, pattern.search(file_name).groups())

                    file_count += 1
                    if file_count % 1000 == 0:
                        print('Copy Image Count: {}'.format(file_count))

                    src_file_name = os.path.join(file_dir, 'F{:04d}B00.jpg'.format(fr_num + 1))

                    src_file_path = os.path.join(args.image_dir, src_file_name)
                    if not os.path.exists(src_file_path):
                        continue

                    dst_file_name = os.path.join(file_dir, 'F{:04d}B00.jpg'.format(fr_num))
                    dst_file_path = os.path.join(out_img_dir, dst_file_name)

                    out_path = os.path.dirname(dst_file_path)
                    if not os.path.exists(out_path):
                        os.makedirs(out_path)
                    shutil.copy(src_file_path, dst_file_path)

                    out_info_str = '{} {}\n'.format(dst_file_name, bbox)
                    out_file.write(out_info_str)

            print('Copy Image Count: {}'.format(file_count))

    out_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {} -Rf'.format(out_img_dir))
    print('finish!')


if __name__ == '__main__':
    main()

